from typing import Dict, List, Optional

from ..instances.BaseOutputFormatter import BaseOutputFormatter
from ..format_utils import *
from ..instances.MainTaskFormatter import MainTaskFormatter


class RelTypeExtractFormatter(MainTaskFormatter):
    """
    Task: Given sentence, extract the relation in the sentence
    """

    def __init__(self, set_option=True, **kwargs) -> None:
        super().__init__(**kwargs)
        self.set_option = set_option

    def format_options(self, options: List[str]) -> str:
        options_joined = ", ".join(options)

        if len(options) > 1:
            options_joined += "."

        return Preffix.OPTIONS.value + options_joined

    def format_data_sample(
        self,
        data_item,
        instruction: str,
        options: Optional[List[str]] = None,
    ) -> Dict:
        sentence = data_item["sentence"]
        relations = data_item["relations"]

        instruction = add_instruction_prefix(instruction)

        options_joined = self.format_options(options) if self.set_option else ""

        sentence = add_sentence_prefix(sentence)

        # answer_format = TypeAnswerFormatter.from_relations(relations, )
        answer_format = self.format_type_answer(relations)

        return {
            "instruction": instruction + " " + options_joined,
            "context": sentence + " " + Preffix.ANSWER.value,
            "labels": answer_format,
        }


def main():
    sample = {
        "sentence": "John Wilkes Booth , who assassinated President Lincoln.",
        "relations": [
            {
                "head": {
                    "name": "John Wilkes Booth",
                    "type": "People",
                    "pos": [0, 17],
                },
                "type": "Kill",
                "tail": {
                    "name": "President Lincoln",
                    "type": "People",
                    "pos": [37, 54],
                },
            }
        ],
        "instruction": "please extract all relation type from the input sentence",
    }

    formatter = RelTypeExtractFormatter()
    print(
        formatter.format_data_sample(
            sample,
            instruction="please extract all relation type from the input sentence",
            options=[
                "Organization based in",
                "Located in",
                "Live in",
                "Work for",
                "Kill",
            ],
        )
    )


if __name__ == "__main__":
    main()
